#!/bin/bash 


__FileName__=""
__Author__="SuoSuo"
__Time__="/2024/04/26"
__Version__="V1.0"




# 显示程序的版本号和更新日期
version () {
	printf "\033[34m
========================================================================= 
*   ${PROGRAM_NAME} ${PROGRAM_VERSION}
*   ${AUTHOR_NAME}          ${AUTHOR_NUMBER}
=========================================================================
Release Note:
    >>V1.0 20240428 first Commit, 可正常执行测试; 
---------------------------------------------------------------------------
    \033[0m\n"
	exit 0
}


# 手动修改值
# 指定 SPDK 和 FIO 源码的安装路径
GIT_SPDK_DIR="/suosuo/spdk-home/spdk"
GIT_FIO_DIR="/suosuo/spdk-home/fio"
# 指定内存大页分配总内存大小，单位: MB
HUGE_MEM_SIZE="10240"



print () {
#    printf "$(gdate): "
    if [ "${1}" = "prompt" ]; then
        printf "\033[32m${2}\033[0m\n"
    elif [ "${1}" = "info" ]; then
        printf "\033[34m${2}\033[0m\n"
    elif [ "${1}" = "error" ]; then
        printf "\033[31m${2}\033[0m\n"
        exit 1 
    else
        printf "${2}\n"
    fi
}

setup_spdk() {
    spdk_root_dir=$1
    white_list="$2"
    huge_mem=$3

    echo ${white_list}

    if [ -z "${huge_mem}" ]
    then
        huge_mem=1024
    fi
    # 添加 DSA 设备 
    dsa_bdf_list="$(ls /sys/bus/pci/drivers/idxd | grep 0000: | xargs)"
    # PCI_ALLOWED="${white_list}" ${spdk_dir}/scripts/setup.sh
    # HUGEMEM=${huge_mem} PCI_ALLOWED="${white_list}" ${spdk_root_dir}/scripts/setup.sh > /dev/null 2>&1
    # HUGE_EVEN_ALLOC="yes"
    HUGE_EVEN_ALLOC="yes" HUGEMEM=${huge_mem} PCI_ALLOWED="${white_list} ${dsa_bdf_list}" ${spdk_root_dir}/scripts/setup.sh
}

reset_spdk() {
    spdk_root_dir=$1
    white_list=$2

    if [ ! -z "${white_list}" ]; then
        PCI_WHITELIST="${white_list}" ${spdk_root_dir}/scripts/setup.sh reset > /dev/null 2>&1

        # reset nvme drive names to original order
        for bdf in ${white_list}; do
            echo ${bdf} > /sys/bus/pci/drivers/nvme/unbind
        done

        for bdf in ${white_list}; do
            echo ${bdf} > /sys/bus/pci/drivers/nvme/bind
        done
    fi
}

# 安全擦除
sercure_erase () {
    for disk_name in "${NVME_NAME_LIST[@]}";do
        nvme format /dev/"$disk_name" -s 1
    done
}

# 工具check
tool_check () {
    ## 系统工具及fio
    for tool_name in $NESS_TOOL_LIST;do
        
        if ! which "$tool_name" &> /dev/null;then
            print "error" "$tool_name not install,please cheack !"
        fi
    done
    if kill -0 $(pgrep sar) > /dev/null 2>&1;then
        print "error" "Please stop the sar process and re-run the script !"
    fi
}


# 收集 CPU 压力日志 
cpu_collect () {
    echo "******开始收集cpu_collect-$1-日志"
    # iostat -mtx 1 > "$tool_path"/iostat_"$1" 2>&1 &
    # iostat_id=$(pgrep iostat)
    sar -P ALL 1 >  "$tool_path"/cpu_lat_"$1" 2>&1 &
    SAR_PID=$(pgrep sar)
}

# 各个盘的信息收集 
disk_log_collect () {
    # nvme info
    nvme list > "$tool_path"/nvme_info
    lsblk -I 259,8 -d -o "NAME,Serial,model,size,state,sched,type,rev,path,mode,vendor" 
    # smart
    echo "" > "$tool_path"/smart_info_"$1" 2>&1
    for disk_name in "${NVME_NAME_LIST[@]}";do
        nvme smart-log /dev/"$disk_name" >> "$tool_path"/smart_info_"$1" 2>&1
    done
}

# 系统日志信息收集
sys_log_collect () {
    # BMC
    ipmitool sel list > "$tool_path"/bmc_"$1"
    # dmesg
    if [ "$1" != "basic" ];then
        dmesg -TL > "$tool_path"/dmesg_"$1"
    fi
    ## 第一次清理及收集完成后进行清理
    dmesg -C

}

# dmesg 信息检查 
dmesg_check () {
    result=$(grep -iE "$dmesg_check_flags" "$tool_path"/dmesg_"$1")
    if [[ $result = "" ]];then
        echo "dmesg_$1 check pass !" >> "$tool_path"/nvme_stress_all_log
    else
        echo "dmesg_$1 check fail !" >> "$tool_path"/nvme_stress_all_log
        echo "$result" >> "$tool_path"/nvme_stress_all_log
    fi
}

# 系统功耗信息收集 
power_diss_collect () {
    echo  

}

# speed_collect flags -> $1 flags -> $2
nvme_speed_collect () {
    mkdir -p "$tool_path"/speed_info_"$2"
    cat /dev/null > "$tool_path"/speed_info_"$2"/lspci_speed_"$1" 2>&1
    for nvme_bus_id in "${NVME_BDF_LIST[@]}";do
        echo "$nvme_bus_id" >> "$tool_path"/speed_info_"$2"/lspci_speed_"$1" 2>&1
        lspci -vvvs "$nvme_bus_id" | grep -i speed | grep -vi lnkctl >> "$tool_path"/speed_info_"$2"/lspci_speed_"$1" 2>&1
    done
}

# speed_check datetime -> $1
speed_check () {
    lspci_speed_check=$(diff "$tool_path"/speed_info_basic/lspci_speed_basic "$tool_path"/speed_info_"$2"/lspci_speed_"$1" 2>&1)
    echo "$1" >> "$tool_path"/nvme_stress_all_log
    
    if [[ $lspci_speed_check != "" ]];then
        echo "lspci_speed_check fail,please check !" >> "$tool_path"/nvme_stress_all_log
        echo "$lspci_speed_check" >> "$tool_path"/nvme_stress_all_log
    else
        echo "lspci_speed_check pass !" >> "$tool_path"/nvme_stress_all_log
    fi
    echo "" >> "$tool_path"/nvme_stress_all_log
}


create_fio_config () {

    fio_config_name="spdk_suosuo.fio"
    cat > $fio_config_name << EOF 
[global]
percentile_list=1:5:10:25:50:75:90:95:99:99.5:99.9:99.99:99.999:99.9999:99.99999:99.999999
ioengine=spdk
direct=1
thread
group_reporting
time_based
verify=0
ramp_time=30

[job]
rw=randread
bs=4k
numjobs=2
iodepth=256
runtime=720
EOF


}
export random_map_opts=${random_map_opts-" --norandommap=1  --randrepeat=0 "}

parase_nvme_info () {
    # 获取 所有的 NVME 盘符
    NVME_NAME_LIST=($(lsblk -I 259 -d -n -o 'NAME' | xargs))
    if [ "${#NVME_NAME_LIST[@]}" -eq 0 ];then
        print "prompt" "This System not exists NVME Device......"
        exit 1
    fi
    # 根据盘符获取其 BDF 号 
    for nvme_name in "${NVME_NAME_LIST[@]}"
    do 
        bdf_id=$(realpath /sys/class/block/${nvme_name// /} | awk -F'/' '{print $6}' | xargs)
        NVME_BDF_LIST=(${NVME_BDF_LIST[@]} ${bdf_id})
    done
    # 根据所有 NVME 盘的BDF 号获取其  node ID 号 
    for bdf_id in "${NVME_BDF_LIST[@]}"
    do
        node_id=$(lspci -nvms ${bdf_id} | awk '/^NUMANode/{print $NF}' | xargs)
        NVME_NODE_ID_LIST=(${NVME_NODE_ID_LIST[@]} ${node_id})
    done  
}

fio_exe () {
    ii=0
    fio_pid_list=""
    printf "\033[31m\n ****** Start Fio Testing Using SPDK ****** \033[0m\n"
    echo -e "\033[31m On Drives: [" "${NVME_NAME_LIST[@]}" "]\033[0m\n"
    # 创建 fio 日志文件夹
    mkdir -p "${SPDK_LOG_DIR_FIO}"
    # 配置 Fio 随机IO负载
    random_map_opts=${random_map_opts-" --norandommap=1  --randrepeat=0 "}

    for bdf in "${NVME_BDF_LIST[@]}"
    do
        fio_filename_format="trtype=PCIe traddr=%s ns=1"
        fio_node_id="--numa_cpu_nodes=${NVME_NODE_ID_LIST[${ii}]}"
        fio_output_file="--output=${SPDK_LOG_DIR_FIO}/${NVME_NAME_LIST[${ii}]}_4K_8_256_randread.log"
        
        # 可执行的命令 :  IOPS 可达到 1760K
        LD_PRELOAD=${LD_PRELOAD_PATH}  \
            ${FIO_SH_PATH} --filename="$(printf "${fio_filename_format}" ${bdf//:/.})" \
            ${fio_node_id} \
            ${random_map_opts} \
            ${fio_output_file} spdk_suosuo.fio &
        fio_pid_list="${fio_pid_list} $!"
        let ii++
    done
    # 等待fio执行完成
    # child_pids=$(pgrep -P $$ -l | grep -w "fio" | awk '{print $1}' | xargs)
    # wait $child_pids
    wait ${fio_pid_list}
    sync
    printf "\n\033[31m Randread 4 2 256 测试执行结束 \033[0m\n\n"
}

spdk_main () {

    # 创建 Fio 文件 
    create_fio_config    
    # 提取所有盘的信息 
    parase_nvme_info

    echo "NVME_NAME_LIST =" "${NVME_NAME_LIST[@]}"
    echo "NVME_BDF_LIST =" "${NVME_BDF_LIST[@]}"
    echo "NVME_NODE_ID_LIST =" "${NVME_NODE_ID_LIST[@]}"

    # 预配置 SPDK 
    setup_spdk "${GIT_SPDK_DIR}" "${NVME_BDF_LIST[*]}" "${HUGE_MEM_SIZE}"
    if [ $? -ne 0 ]
    then
        echo "setup spdk failed, revert ..."
        reset_spdk "${GIT_SPDK_DIR}" "${NVME_BDF_LIST[*]}"
        echo "revert done"
        exit 1
    fi

    # 开始执行测试
    fio_exe

    # 解除 SPDK 配置 
    reset_spdk "${GIT_SPDK_DIR}" "${NVME_BDF_LIST[*]}"

}


# 当前时间，格式是"年.月.日.时.分.秒"
date_time=$(date +%Y.%m.%d.%H.%M.%S)
# 文件名称
PROGRAM_NAME=$(basename $0)
# 作者名字
AUTHOR_NAME="SuoSuo"
# 作者工号
AUTHOR_NUMBER="47372"
# 版本
VERSION="V1.1.0"
# 首地址
BASE_DIR=$(pwd)

path=$(dirname $0)
[[ "$path" = "." ]] && { path=$(pwd) ; BASE_DIR=$path; cd $path; }
# test "${BASE_DIR}" != "/home" && { print "info" "must path is /home;" ; exit 1; }




# NVME 所有盘的盘符
NVME_NAME_LIST=()
# NVME 盘的所有 BDF 号
NVME_BDF_LIST=()
# NVME 盘的所有 NODE 号 
NVME_NODE_ID_LIST=()

# FIO 可执行文件的全路径 
FIO_SH_PATH="${GIT_FIO_DIR}/fio"
# spdk_nvme 的全路径
LD_PRELOAD_PATH="${GIT_SPDK_DIR}/build/fio/spdk_nvme"

[ ! -e "${FIO_SH_PATH}" ] && print "error" "The file or directory does not exist: ${FIO_SH_PATH}"
[ ! -e "${LD_PRELOAD_PATH}" ]  && print "error" "The file or directory does not exist: ${LD_PRELOAD_PATH}"

# 指定总的日志文件夹
SPDK_LOG_DIR="fio2spdk_log_$(date +%Y%m%d_%H%M%S)"
# 指定 Fio 执行测试结果文件夹
SPDK_LOG_DIR_FIO="${SPDK_LOG_DIR}/fio_randread_res"
# 指定 Fio 执行测试结果
SPDK_LOG_DIR_FIO_MSG="${SPDK_LOG_DIR_FIO}/fio_randread_log"
# 指定环境收集日志文件夹
SPDK_LOG_DIR_MSG="${SPDK_LOG_DIR}/collect_message"
# 需要安装的工具列表
NESS_TOOL_LIST="lsblk nvme sar fio ipmitool smartctl"



if [ -e "${SPDK_LOG_DIR}" ]; then
    print "prompt" " Dir : ${SPDK_LOG_DIR} already exists .... \n please remove old file...\n"
    exit 1
else 
    mkdir ${SPDK_LOG_DIR}
fi


spdk_main